#ifndef CLASS_AG
#define CLASS_AG

#include <stdio.h>
#include <smmintrin.h>

#include "TreeForceFinder.h"
#include "Vector.h"


#define SQR(x) ((x)*(x))


class AG {
	public:
				
	int nr, nz;
				
	REAL r1, r2;
	REAL z1, z2;
				
	int nphi, nsign;
		
	REAL ErrTolTheta, softening;
	
	REAL G;
	
	typedef long int lint;
	
	// derivatives
	INT r1i,r2i,z1i,z2i,dri,dzi, n;
	
	class A {  
		public:
			
		REAL r,z,pot;
		
		A():r(0),z(0),pot(0){};
		
		void print(){ 
			printf("{%g, %g}, %g\n",r,z, pot); 
		}
		
	} *a;

	union FloatUintUnion { float f; unsigned int i;	};
	union DoubleLuintUnion { double f; long unsigned int i; };
	
	//typedef long unsigned int luint;
	
	AG() :
	G(1),
	a(0)
	{ 
		r1 = 1; 
		nr = 22*100; 
		r2 = pow(2,22);
		
		z1 = 1; 
		nz = 22*100; 
		z2 = pow(2,22);
		
		nsign = 1; 
		nphi  = 1;
		
		ErrTolTheta = 0.35; 
		softening = 0.1;
		
		update_derivatives(); 
	}
				
				
	void update_derivatives() {
					
		//inv_dr = nr/(r2 - r1); 
		//inv_dz = nz/(z2 - z1);

		r1i = *(INT*)&r1;
		r2i = *(INT*)&r2;
		
		z1i = *(INT*)&z1;		
		z2i = *(INT*)&z2;
		
		dri = (r2i-r1i)/nr;
		dzi = (z2i-z1i)/nz;
		
		inv_dr = REAL(nr)/(r2i-r1i); 
		inv_dz = REAL(nz)/(z2i-z1i);
		
		r1jj = _mm_set1_ps(r1);
		r2jj = _mm_set1_ps(r2);
		
		z1jj = _mm_set1_ps(z1);
		z2jj = _mm_set1_ps(z2);
		
		r1ijj = _mm_set1_epi32(r1i);
		r2ijj = _mm_set1_epi32(r2i);
		
		z1ijj = _mm_set1_epi32(z1i);
		z2ijj = _mm_set1_epi32(z2i);
		
		drijj = _mm_set1_epi32(dri);
		dzijj = _mm_set1_epi32(dzi);
		
		inv_drjj = _mm_set1_ps(inv_dr);
		inv_dzjj = _mm_set1_ps(inv_dz);
		
		onejj = _mm_set1_epi32(1);
		nr1jj = _mm_set1_epi32(nr+1);
		
		nrjj = _mm_set1_epi32(nr);
		nzjj = _mm_set1_epi32(nz);
		
		n = (nr+1)*(nz+1);
		
		free();
		a = new A[n];

	}
				
				
	void calc(P *p, int n) {
					
		printf ("AG::calc()::sizeof(a) = %g [MB]\n", 
				  sizeof(A)*nr*nz/1024.0/1024 );
					
		TreeForceFinder tff;
		tff.G = G;
		tff.ErrTolTheta = ErrTolTheta;
		tff.softening = softening;

		tff.NumPart = n;
		tff.force_treeallocate(1.5*n, n);
		tff.force_treebuild(p);
					
		int sign[2] = {+1,-1};
					
					
		for (int iz=0; iz<=nz; iz++)
		for (int ir=0; ir<=nr; ir++) {
						
		
			INT ri = r1i + ir*dri;
			INT zi = z1i + iz*dzi;

			REAL r = -r1 + *(REAL*)&ri;
			REAL z = -z1 + *(REAL*)&zi;
			
			A ai;
			
			REAL acci[3], posi[3], pot(0);
						
			for (int iphi=0; iphi<nphi; iphi++)
			for (int isign =0; isign<nsign; isign++) {
					
				REAL alpha = (REAL)iphi/nphi*2*M_PI;
				REAL u[] = { cos(alpha), sin(alpha), sign[isign] };
							
				posi[0] = r*u[0];
				posi[1] = r*u[1];
				posi[2] = z*u[2];
							
				tff.force_treeevaluate(p, posi, acci, pot);
							
				ai.r+= acci[0]*u[0] + acci[1]*u[1]; 
				ai.z+= acci[2]*u[2];
				ai.pot+=pot;
			}
			
			if (ir==0) ai.r = 0;
			if (iz==0) ai.z = 0;
			
			a[iz*(nr+1) + ir].r+=   ai.r/(nphi*nsign);
			a[iz*(nr+1) + ir].z+=   ai.z/(nphi*nsign);
			a[iz*(nr+1) + ir].pot+= ai.pot/(nphi*nsign);

			printf("%i %i\n", ir,iz);
		}
					
		tff.force_treefree();

	}
	
	
	inline A get_acc(REAL r, REAL z) {
		
		REAL rr1(r+r1);
		REAL zz1(z+z1);
		
		REAL ird = (*(INT*)&rr1-r1i) * inv_dr; 
		REAL izd = (*(INT*)&zz1-z1i) * inv_dz; 

		int ir = ird;
		int iz = izd;
		
		return get_acc(rr1,zz1, ird,izd, ir,iz);
		
	}
	
	
	inline A get_acc_ps(__m128 &rjj, __m128 &zjj, __m128 &accrjj, __m128 &acczjj) {
		
		//// REAL rr1(r+r1); ////
		//// REAL zz1(z+z1); ////
		
		__m128 rr1jj  = _mm_add_ps(rjj, r1jj); // ok
		__m128 zz1jj  = _mm_add_ps(zjj, z1jj); // ok

		
		//// REAL ird = (*(INT*)&rr1-r1i) * inv_dr; ////
		//// REAL izd = (*(INT*)&zz1-z1i) * inv_dz; ////
		
		__m128 irdjj = _mm_mul_ps(_mm_cvtepi32_ps(_mm_sub_epi32(_mm_castps_si128(rr1jj), r1ijj)), inv_drjj); // ok
		__m128 izdjj = _mm_mul_ps(_mm_cvtepi32_ps(_mm_sub_epi32(_mm_castps_si128(zz1jj), z1ijj)), inv_dzjj); // ok

		
		//// int ir = ird; ////
		//// int iz = izd; ////
		
		__m128i irjj = _mm_cvttps_epi32(irdjj); // ok
		__m128i izjj = _mm_cvttps_epi32(izdjj); // ok

		
		//// int i = iz*(nr+1)+ir; ////
		
		__m128i maskjj =  _mm_and_si128(_mm_cmplt_epi32(irjj, nrjj), _mm_cmplt_epi32(izjj, nzjj));
		irjj = _mm_and_si128(maskjj, irjj);
		izjj = _mm_and_si128(maskjj, izjj);
		
		
		//// A &ai  = a[i];  ////
		//// A &air = a[i+1];  ////
		//// A &aiz = a[i+nr+1];  ///
		
		__m128i ijj = _mm_add_epi32( _mm_mullo_epi32(izjj, nr1jj), irjj); // ok
		__m128i i1jj = _mm_add_epi32(ijj, onejj);  // ok
		__m128i inr1jj = _mm_add_epi32(ijj, nr1jj);  // ok
	
		INT i[4],i1[4],inr1[4];
		
		_mm_store_si128((__m128i*)i, ijj);// ok
		_mm_store_si128((__m128i*)i1, i1jj); // ok
		_mm_store_si128((__m128i*)inr1, inr1jj); // ok
		
		__m128 airjj = _mm_setr_ps( a[i[0]].r, a[i[1]].r, a[i[2]].r, a[i[3]].r ); // ok
		__m128 aizjj = _mm_setr_ps( a[i[0]].z, a[i[1]].z, a[i[2]].z, a[i[3]].z ); // ok
		
		__m128 airrjj = _mm_setr_ps( a[i1[0]].r, a[i1[1]].r, a[i1[2]].r, a[i1[3]].r ); // ok
		__m128 airzjj = _mm_setr_ps( a[i1[0]].z, a[i1[1]].z, a[i1[2]].z, a[i1[3]].z );// ok
		
		__m128 aizrjj = _mm_setr_ps( a[inr1[0]].r, a[inr1[1]].r, a[inr1[2]].r, a[inr1[3]].r );// ok
		__m128 aizzjj = _mm_setr_ps( a[inr1[0]].z, a[inr1[1]].z, a[inr1[2]].z, a[inr1[3]].z );// ok
	
		
		//// REAL drn = ird-ir;  ////
		//// REAL dzn = izd-iz;  ////
		//// accrjj = airjj + (airrjj-airjj)*drnjj + (aizrjj-airjj)*dznjj;  ////
		//// acczjj = aizjj + (airzjj-aizjj)*drnjj + (aizzjj-aizjj)*dznjj;  ////

		__m128 drnjj  = _mm_sub_ps(irdjj, _mm_cvtepi32_ps(irjj)); // ok
		__m128 dznjj  = _mm_sub_ps(izdjj, _mm_cvtepi32_ps(izjj));// ok
		
		accrjj = _mm_add_ps( _mm_add_ps(airjj, _mm_mul_ps(_mm_sub_ps(airrjj, airjj), drnjj)), _mm_mul_ps(_mm_sub_ps(aizrjj, airjj), dznjj) );// ok
		acczjj = _mm_add_ps( _mm_add_ps(aizjj, _mm_mul_ps(_mm_sub_ps(airzjj, aizjj), drnjj)), _mm_mul_ps(_mm_sub_ps(aizzjj, aizjj), dznjj) );// ok
		/**/
		
		accrjj = _mm_and_ps(_mm_castsi128_ps(maskjj), accrjj); // ok
		acczjj = _mm_and_ps(_mm_castsi128_ps(maskjj), acczjj); // ok
		
	}
	
	
	
	
	
	
	
	
	
	
	inline A get_acc(REAL&r,REAL&z, REAL&ird,REAL&izd, int&ir,int&iz) {

		A arz;
		
		if (ir<nr && iz<nz ) {
		
			int i = iz*(nr+1)+ir;
			
			
			
			A &ai  = a[i];
			A &air = a[i+1];
			A &aiz = a[i+nr+1];
			
			/*
			printf("%7.3f %7.3f  %7.3f %7.3f  %7.3f %7.3f\n", 
					 ai.r, ai.z, air.r, air.z, aiz.r, aiz.z);
				*/	 
			REAL drn = ird-ir;
			REAL dzn = izd-iz;
		
			arz.r = ai.r + (air.r-ai.r)*drn + (aiz.r-ai.r)*dzn;
			arz.z = ai.z + (air.z-ai.z)*drn + (aiz.z-ai.z)*dzn;
		
		} 
		
		return arz;
	}
	
	
	inline A get_acc_pot(REAL r, REAL z) {
		
		REAL rr1(r+r1);
		REAL zz1(z+z1);
		
		REAL ird = (*(INT*)&rr1-r1i) * inv_dr;
		REAL izd = (*(INT*)&zz1-z1i) * inv_dz;

		int ir = ird;
		int iz = izd;
		
		return get_acc_pot(rr1,zz1, ird,izd, ir,iz);
		
	}
	
	bool isInGrid(REAL *pos) {
	
		REAL r = sqrt(SQR(pos[0])+SQR(pos[1]));
		REAL z = fabs(pos[2]);
		
		REAL rr1(r+r1);
		REAL zz1(z+z1);
		
		int ir = (*(INT*)&rr1-r1i) * inv_dr;
		int iz = (*(INT*)&zz1-z1i) * inv_dz;
		
		return  ir<=nr && iz<=nz  ? true : false; 
		
	}

	A get_acc_pot(REAL&r,REAL&z, REAL&ird,REAL&izd, int&ir,int&iz) {
		
		A arz;
		
		if (ir<nr && iz<nz ) {
		
			int i = iz*(nr+1)+ir;
			A &ai  = a[i];
			A &air = a[i+1];
			A &aiz = a[i+nr+1];
					
			REAL drn = ird-ir;
			REAL dzn = izd-iz;
		
			arz.r = ai.r + (air.r-ai.r)*drn + (aiz.r-ai.r)*dzn;
			arz.z = ai.z + (air.z-ai.z)*drn + (aiz.z-ai.z)*dzn;
			arz.pot = ai.pot + (air.pot-ai.pot)*drn + (aiz.pot-ai.pot)*dzn;
		}

		return arz;
		
	}

	
	A get_acc_pot(REAL *pos) {
		
		REAL r2 = SQR(pos[0])+SQR(pos[1]);
		REAL r = sqrt(r2);
		REAL z = fabs(pos[2]);
		
		REAL rr1(r+r1);
		REAL zz1(z+z1);
		
		REAL ird = (*(INT*)&rr1-r1i) * inv_dr;
		REAL izd = (*(INT*)&zz1-z1i) * inv_dz;

		int ir = ird;
		int iz = izd;
		
		A arz;
		
		if (ir<nr && iz<nz ) {
		
			int i = iz*(nr+1)+ir;
			A &ai  = a[i];
			A &air = a[i+1];
			A &aiz = a[i+nr+1];
					
			REAL drn = ird-ir;
			REAL dzn = izd-iz;
		
			arz.r = ai.r + (air.r-ai.r)*drn + (aiz.r-ai.r)*dzn;
			arz.z = ai.z + (air.z-ai.z)*drn + (aiz.z-ai.z)*dzn;
			arz.pot = ai.pot + (air.pot-ai.pot)*drn + (aiz.pot-ai.pot)*dzn;
		}

		return arz;
		
	}
	
	
	
	
	void save(const char *filename) {
					
		printf("AG::save('%s')...", filename); fflush(stdout);
		FILE *fp = fopen(filename, "wb");
		fwrite(this, sizeof(*this), 1, fp);
		fwrite(a, sizeof(*a), (nr+1)*(nz+1), fp);
		fclose(fp);
		printf("done\n");

	}
				
	bool load(const char *filename) {
					
		printf("AG::load('%s')...", filename); fflush(stdout);
		if (FILE *fp = fopen(filename, "rb")) {
		
			free();
			
			fread(this, sizeof(*this), 1, fp);
					 
			a = new A[(nr+1)*(nz+1)];
			fread(a, sizeof(*a), (nr+1)*(nz+1), fp);

			fclose(fp);
			printf("done\n");
			return true;
						
		} else {
						
			printf("failed\n");
			return false;
		}
					
	}
	
	
	void save_txt(const char *filename){
		
		printf("AG::save_txt('%s')...", filename); fflush(stdout);
		
		FILE *fp = fopen(filename, "w");
		
		for (int iz=0; iz<=nz; iz++)
		for (int ir=0; ir<=nr; ir++) {
			
			INT ri = r1i + ir*dri;
			INT zi = z1i + iz*dzi;
			
			REAL r = -r1 + *(REAL*)&ri;
			REAL z = -z1 + *(REAL*)&zi;
			
			A &ai = a[iz*(nr+1) + ir];
			fprintf(fp, "%.3e\t%.3e\t%.3e\t%.3e\t%.2e\n", 
						r, z, ai.r, ai.z, ai.pot);
			
		}
		
		fclose(fp);
		
		printf("done\n");
	}
	
	void free() {
		if (a!=0) delete [] a;
		a = 0;
	}
				
	void print(){
					
		printf("AG::print()\n");
		printf(" ├ r1 = %g %u\n", r1, r1i);
		printf(" ├ nr = %i\n", nr);
		printf(" ├ r2 = %g %u\n", r2, r2i);
		printf(" ├ z1 = %g %u\n", z1, z1i);
		printf(" ├ nz = %i\n", nz);
		printf(" ├ z2 = %g %u\n", z2, z2i);
		printf(" ├ nsign = %i\n", nsign);
		printf(" ├ nphi = %i\n", nphi);
		printf(" └ G = %g\n\n", G);

	}
				
	private:
		REAL inv_dr, inv_dz;
		
		__m128 r1jj, r2jj;
		__m128 z1jj, z2jj;
		
		__m128i r1ijj, r2ijj;
		__m128i z1ijj, z2ijj;
		__m128i drijj, dzijj;
		
		__m128 inv_drjj, inv_dzjj;
		__m128i nrjj, nzjj;
		__m128i onejj, nr1jj;
};

#endif	
